import sys

sys.path.append("../../../")

from common.python.common import *
from common.python.ipe_param import *


class IpeParamDetect(IpeParamBase):
    def __init__(self, model_width, model_height):
        super().__init__(model_width, model_height)
        # 图片resize参数
        self.m_iResizeWidth = 0
        self.m_iResizeHeight = 0
        # 图片填充参数
        self.m_iPadRight = 0
        self.m_iPadBottom = 0

    def calc_param(self, stream, input_data, output_data):
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescIn))
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescOut))
        error_check(sdk.lyn_ipe_reset_config_desc(self.m_oConfigDesc))
        error_check(
            sdk.lyn_ipe_set_input_pic_desc(
                self.m_oPicDescIn,
                input_data,
                self.m_iImgWidth,
                self.m_iImgHeight,
                self.m_format,
            )
        )
        error_check(sdk.lyn_ipe_set_output_pic_data(self.m_oPicDescOut, output_data))
        error_check(
            sdk.lyn_ipe_set_c2c_config(
                self.m_oConfigDesc, sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24, 0
            )
        )
        error_check(
            sdk.lyn_ipe_set_resize_config(
                self.m_oConfigDesc, self.m_iResizeWidth, self.m_iResizeHeight
            )
        )
        error_check(
            sdk.lyn_ipe_set_pad_config(
                self.m_oConfigDesc,
                self.m_iPadBottom,
                self.m_iPadRight,
                self.m_iPadBottom,
                self.m_iPadRight,
                114,
                114,
                114,
            )
        )
        error_check(
            sdk.lyn_ipe_cal_output_pic_desc(
                self.m_oPicDescOut, self.m_oPicDescIn, self.m_oConfigDesc, 2
            )
        )
        error_check(
            sdk.lyn_ipe_process_async(
                stream, self.m_oPicDescIn, self.m_oPicDescOut, self.m_oConfigDesc
            )
        )

    def set_img_info(self, width, height, format):
        scale = max(width / self.m_iModelWidth, height / self.m_iModelHeight)
        self.m_iResizeWidth = int(width / scale) & ~1
        self.m_iResizeHeight = int(height / scale) & ~1

        if self.m_iResizeWidth > 1920:
            self.m_iResizeWidth = 1920
        if self.m_iResizeHeight > 1080:
            self.m_iResizeHeight = 1080

        self.m_iPadRight = (self.m_iModelWidth - self.m_iResizeWidth) // 2
        self.m_iPadBottom = (self.m_iModelHeight - self.m_iResizeHeight) // 2

        super().set_img_info(width, height, format)


class IpeParamRecog(IpeParamBase):
    def __init__(self, model_width, model_height):
        super().__init__(model_width, model_height)
        # 图片resize参数
        self.resize_w = 0
        self.resize_h = 0
        # 图片填充参数
        self.pad_b = 0
        self.pad_r = 0

    def calc_param(self, stream, input_data, output_data):
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescIn))
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescOut))
        error_check(sdk.lyn_ipe_reset_config_desc(self.m_oConfigDesc))
        error_check(
            sdk.lyn_ipe_set_input_pic_desc(
                self.m_oPicDescIn,
                input_data,
                self.m_iImgWidth,
                self.m_iImgHeight,
                self.m_format,
            )
        )
        error_check(sdk.lyn_ipe_set_output_pic_data(self.m_oPicDescOut, output_data))
        error_check(
            sdk.lyn_ipe_set_c2c_config(
                self.m_oConfigDesc, sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24, 0
            )
        )
        error_check(
            sdk.lyn_ipe_set_resize_config(
                self.m_oConfigDesc, self.resize_w, self.resize_h
            )
        )
        error_check(
            sdk.lyn_ipe_set_pad_config(
                self.m_oConfigDesc,
                0,
                self.pad_r,
                self.pad_b,
                0,
                114,
                114,
                114,
            )
        )
        error_check(
            sdk.lyn_ipe_cal_output_pic_desc(
                self.m_oPicDescOut, self.m_oPicDescIn, self.m_oConfigDesc, 2
            )
        )
        error_check(
            sdk.lyn_ipe_process_async(
                stream, self.m_oPicDescIn, self.m_oPicDescOut, self.m_oConfigDesc
            )
        )

    def set_img_info(self, width, height, format):
        ratio = width / height
        new_w = ratio * self.m_iModelHeight
        self.resize_w = int(min(new_w, self.m_iModelWidth))
        self.resize_h = int(self.m_iModelHeight)

        self.pad_b = self.m_iModelHeight - self.resize_h
        self.pad_r = self.m_iModelWidth - self.resize_w

        super().set_img_info(width, height, format)


class IpeParamW600k_mbf(IpeParamBase):
    def __init__(self, model_width, model_height):
        super().__init__(model_width, model_height)

    def calc_param(self, stream, input_data, output_data):
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescIn))
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescOut))
        error_check(sdk.lyn_ipe_reset_config_desc(self.m_oConfigDesc))
        error_check(
            sdk.lyn_ipe_set_input_pic_desc(
                self.m_oPicDescIn,
                input_data,
                self.m_iImgWidth,
                self.m_iImgHeight,
                self.m_format,
            )
        )
        error_check(sdk.lyn_ipe_set_output_pic_data(self.m_oPicDescOut, output_data))
        error_check(
            sdk.lyn_ipe_set_c2c_config(
                self.m_oConfigDesc, sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24, 0
            )
        )
        error_check(
            sdk.lyn_ipe_set_resize_config(
                self.m_oConfigDesc, self.m_iModelWidth, self.m_iModelHeight
            )
        )
        error_check(
            sdk.lyn_ipe_cal_output_pic_desc(
                self.m_oPicDescOut, self.m_oPicDescIn, self.m_oConfigDesc, 2
            )
        )
        error_check(
            sdk.lyn_ipe_process_async(
                stream, self.m_oPicDescIn, self.m_oPicDescOut, self.m_oConfigDesc
            )
        )


class IpeParamCrop(IpeParamBase):
    def __init__(self, x, y, w, h):
        super().__init__(0, 0)
        self.m_x = x
        self.m_y = y
        self.m_w = w
        self.m_h = h

    def calc_param(self, stream, input_data, output_data):
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescIn))
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescOut))
        error_check(sdk.lyn_ipe_reset_config_desc(self.m_oConfigDesc))
        error_check(
            sdk.lyn_ipe_set_input_pic_desc(
                self.m_oPicDescIn,
                input_data,
                self.m_iImgWidth,
                self.m_iImgHeight,
                self.m_format,
            )
        )
        error_check(sdk.lyn_ipe_set_output_pic_data(self.m_oPicDescOut, output_data))
        error_check(
            sdk.lyn_ipe_set_crop_config(
                self.m_oConfigDesc, self.m_x, self.m_y, self.m_w, self.m_h
            )
        )
        error_check(
            sdk.lyn_ipe_set_c2c_config(
                self.m_oConfigDesc, sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24, 0
            )
        )
        error_check(
            sdk.lyn_ipe_cal_output_pic_desc(
                self.m_oPicDescOut, self.m_oPicDescIn, self.m_oConfigDesc, 0
            )
        )
        error_check(
            sdk.lyn_ipe_process_async(
                stream, self.m_oPicDescIn, self.m_oPicDescOut, self.m_oConfigDesc
            )
        )


class IpeParamAffine(IpeParamBase):
    def __init__(self, iModelWidth, iModelHeight, matrix, xColor, yColor, zColor):
        super().__init__(iModelWidth, iModelHeight)
        self.m_matrix = matrix
        self.m_xColor = xColor
        self.m_yColor = yColor
        self.m_zColor = zColor

    def calc_param(self, stream, input_data, output_data):
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescIn))
        error_check(sdk.lyn_ipe_reset_pic_desc(self.m_oPicDescOut))
        error_check(sdk.lyn_ipe_reset_config_desc(self.m_oConfigDesc))
        error_check(
            sdk.lyn_ipe_set_input_pic_desc(
                self.m_oPicDescIn,
                input_data,
                self.m_iImgWidth,
                self.m_iImgHeight,
                self.m_format,
            )
        )
        error_check(sdk.lyn_ipe_set_output_pic_data(self.m_oPicDescOut, output_data))
        error_check(
            sdk.lyn_ipe_set_output_pic_desc(
                self.m_oPicDescOut,
                output_data,
                self.m_iModelWidth,
                self.m_iModelHeight,
                sdk.lyn_pixel_format_t.LYN_PIX_FMT_RGB24,
            )
        )
        error_check(
            sdk.lyn_ipe_set_affine_config(
                self.m_oConfigDesc,
                self.m_matrix,
                self.m_xColor,
                self.m_yColor,
                self.m_zColor,
            )
        )

        sdk.lyn_ipe_process_async(
            stream, self.m_oPicDescIn, self.m_oPicDescOut, self.m_oConfigDesc
        )
